"""
created: 2023-03-26 14:05
author:qinguoming
description:
1.  这个程序是一个计算模板,用于对3D线性8节点单元组成的结构的静态分析,输出节点位移,输出单元应力,单元应变,节点支反力等结果.
2.  节点自由度为: ux,uy,uz
3.  节点和网格的编号必须从1开始,且连续.
4.  所有set_name,surface_name必须是小写
5.  计算结果输出到mat文件

程序支持的边界条件和加载:
1. 位移边界条件(m):
    定义格式:
        a. (bc_name : str , node_label_list : list[int] , node_dof : int ,displace_value : float)
        b. (bc_name : str , node_set_name : str , node_dof : int ,displace_value : float)
        note: node_dof = 1 or 2 or 3 , 分别代表 ux , uy , uz
        e.g:
            ("DisplacementBc",[1,2,3,...],1,0.0)
2. 节点集中力加载(N):
    定义格式:
        a. ("ConcentratedLoad",load_name : str , node_label_list : list[int] , node_dof : int , force_value : float)
        b. ("ConcentratedLoad",load_name : str , node_set_name : str , node_dof : int , force_value : float)
        e.g:
            ("ConcentratedLoad","load_2000N","nset_name",2,2000.0)];
            ("ConcentratedLoad","load_2000N",[1,2,3,4] ,2,2000.0)];
3. 单元面压力加载(Pa):
    定义格式:
        a. ("PressureLoad" , load_name : str ,surface_set_name : str , pressure_value : float)
        b. ("PressureLoad" , load_name : str ,array_of_element_label_and_faceId : ndarray(m,2), pressure_value : float)
        e.g:
            ("PressureLoad" , "pressure_10Kpa" , "surf-1＂, 10000.0)
            surf_data=np.array([[1,1],[2,3],[15,2],[5,1]...])
            ("PressureLoad" , "pressure_5pa" ,surf_data, 5.0)
4. 体积力加载(N/m^3):
    定义格式:
        a. ("BodyLoad" , load_name : str , el_set_name : str , load_direction : str , body_force_value : float)
        b. ("BodyLoad" , load_name : str , element_label_list : list[int] , load_direction : str , body_force_value : float)
        note : load_direction = 'x' or 'y' or 'z'
        e.g:
            ("BodyLoad" , "gravity_x_10k" , "elset-1＂, 'x' ,10000.0)
            ("BodyLoad" , "gravity_y_5" ,element_label_list, 'y' , 5.0)
5. 面力加载(pa):
    定义格式:
        a. ("SurfaceLoad" , load_name : str ,surface_set_name : str , direction_unit_vector = ndarray (3,) , value : float)
        b. ("SurfaceLoad" , load_name : str ,array_of_element_label_and_faceId : ndarray(m,2) , direction_unit_vector = ndarray (3,) , value : float)
        e.g:
            ("SurfaceLoad" , "pressure_10Kpa" , "surf-1＂, [1.0,0,0] ,10000.0)
            surf_data=np.array([[1,1],[2,3],[15,2],[5,1]...])
            ("SurfaceLoad" , "pressure_5pa" ,surf_data, [1,0,0], 5.0)
6. 温度加载(摄氏度):
    定义格式:
        a. ("ThermalLoad" , load_name : str , el_set_name : str , thermal value : float)
        b. ("ThermalLoad" , load_name : str , element_label_list : list[int] , thermal value : float)
        e.g:
            ("ThermalLoad" , "x" , "elset-1＂,10000.0)
            ("ThermalLoad" , "y" ,element_label_list,  5.0)
"""
import sys
import os
import dill
from EF2D import *
import numpy as np
import scipy as sc
from icecream import ic
from EF2D.tools import *
from EF2D.tools import *
from EF2D.C3D8 import C3D8
from EF2D.Constants import *
from typing import List,Dict
from datetime import datetime
import matplotlib.pyplot as plt
initNumpyPrintOptions()
# -------------定义变量------------------
job_name="job-c3d8-case250407-2"
# 网格文件
# mesh_file="L:\\EF2D\\tests\\data\\Job-F300.inp"
mesh_file="L:/EF2D/tests/AbaqusModelTests/250407/Job250407C3D8COMPLEXMODEL.inp"
# 材料属性(SI-mm单位制)
thermal_expand= 14.4
rho= 1.932e-05
E=77200
nu=0.3
# 约束:支持位移边界条件
def_bc=[("DisplacementBc","Set_fix",1,0.0),
        ("DisplacementBc","Set_fix",2,0.0),
        ("DisplacementBc","Set_fix",3,0.0)]
# 定义载荷(载荷类型,载荷参数)
# def_load=[("ConcentratedLoad","load-1","Set_forceX_100N",1,300.0),]
def_load=[("ConcentratedLoad","load-1","Set_Fy100N",2,-100.0),]
# def_load=[("BodyLoad" ,"body_force-1","body_force_gravity_nz",'z',-500.0)]
# def_load=[ ("SurfaceLoad" , "surface_load-1000pa" , "surface_load_tensile_1000N", [-1.0,0,0] ,1000.0)]
# def_load=[("PressureLoad" , "pressure_500pa" , "pressure_load_500mpa", 5.0)]
# def_load=[ ("ThermalLoad" , "x" , "elset-1＂,10000.0)]

#-////////求解设置//////////////////////
bcMethod=2  # 边界条件处理方法,为1,2,3,4分别是Solver.py的方法
export_result=True  # 是否输出结果
export_path="L:/EF2D/tests/result/"

# -------------导入网格------------------
node_data, element_data, nset_data, elset_data ,surf_data= read_abaqus_inp(mesh_file)
print("网格文件:",mesh_file)
print("节点数:",node_data.shape[0])
print("单元数:",element_data.shape[0])
print("节点集:",nset_data.keys())
print("单元集:",elset_data.keys())
print("surface集合:",surf_data.keys())
print("边界条件:",[b[0] for b in def_bc])
print("载荷有:",[l[1] for l in def_load])
print("边界条件处理方法:",bcMethod)

# -------------建立网格------------------
node_num=node_data.shape[0]
element_num=element_data.shape[0]
flobal_dof_num=node_num*3
# 定义节点对象
node_dict:Dict[int,Node]={}
for i in range(node_num):
    if node_data[i,0] not in node_dict.keys():
        node_dict[node_data[i,0]]=Node(label=node_data[i,0],x=node_data[i,1],y=node_data[i,2],z=node_data[i,3])
# 定义单元对象
element_list:List[C3D8]=[]
for i in range(element_num):
    element_list.append(C3D8(label=element_data[i,0],
                                        node1=node_dict[element_data[i,1]],
                                        node2=node_dict[element_data[i,2]],
                                        node3=node_dict[element_data[i,3]],
                                        node4=node_dict[element_data[i,4]],
                                        node5=node_dict[element_data[i,5]],
                                        node6=node_dict[element_data[i,6]],
                                        node7=node_dict[element_data[i,7]],
                                        node8=node_dict[element_data[i,8]],
                                        E=E,nu=nu,density=rho,expand=thermal_expand))
# 定义载荷和边界
bc_list=[]
for b in def_bc:
    if b[0]=="DisplacementBc":
        if isinstance(b[1],str) and b[1].lower() in nset_data.keys():
            bc_list.append(DisplacementBc(name=b[0],nodeLabels=nset_data[b[1].lower()],dofth=b[2],value=b[3]))
        if (isinstance(b[1],list) or isinstance(b[1],np.ndarray)):
            bc_list.append(DisplacementBc(*b))
load_list=[]
for l in def_load:
    if l[0]=="ConcentratedLoad":
        if isinstance(l[2],str) and l[2].lower() in nset_data.keys():
            load_list.append(ConcentratedLoad(name=l[1],
                                            nodeLabels=nset_data[l[2].lower()],
                                            dofth=l[3],
                                            value=l[4]))
        else:
            load_list.append(ConcentratedLoad(*l[1:]))
    elif l[0]=="PressureLoad":
        if isinstance(l[2],str) and l[2].lower() in surf_data.keys():
            load_list.append(PressureLoad(name=l[1],
                                        ElementLabels=surf_data[l[2].lower()][:,0],
                                        faceIds=surf_data[l[2].lower()][:,1],
                                        value=l[3]))
        else:
            load_list.append(PressureLoad(name=l[1],
                                        ElementLabels=l[2][:,0],
                                        faceIds=l[2][:,1],
                                        value=l[3]))
    elif l[0] == "BodyLoad":
        if isinstance(l[2],str) and l[2].lower() in elset_data.keys():
            load_list.append(BodyLoad(name=l[1],
                                    ElementLabels=elset_data[l[2].lower()],
                                    Direction=l[3],
                                    value=l[4]))
        else:
            load_list.append(BodyLoad(*l[1:]))
    elif l[0] == "SurfaceLoad":
        if isinstance(l[2],str) and l[2].lower() in surf_data.keys():
            load_list.append(SurfaceLoad(name=l[1],
                                        ElementLabels=surf_data[l[2].lower()][:,0],
                                        faceIds=surf_data[l[2].lower()][:,1],
                                        directionVector=l[3],
                                        value=l[4]))
        else:
            load_list.append(SurfaceLoad(name=l[1],
                                        ElementLabels=l[2][:,0],
                                        faceIds=l[2][:,1],
                                        directionVector=l[3],
                                        value=l[4]))
    # ("ThermalLoad" , "y" ,element_label_list,  5.0)
    elif l[0] == "ThermalLoad":
        if isinstance(l[2],str) and l[2].lower() in elset_data.keys():
            load_list.append(ThermalLoad(name=l[1],
                                    ElementLabels=elset_data[l[2].lower()],
                                    temperature=l[4]))
        else:
            load_list.append(BodyLoad(*l[1:]))
# -------------组装刚度矩阵-------------------------
KG=np.zeros((flobal_dof_num,flobal_dof_num))
for e in element_list:
    KG[np.ix_(e.ElemDofIndexs,e.ElemDofIndexs)]+=e.Ke_bbar
# --------------边界条件处理-------------------------
bc_dict={}
for bc in bc_list:
    if isinstance(bc,DisplacementBc):
        for n_label in bc.NodeLabelList:
            dof=get_ithdof_index(n_label,3,bc.Dofth)
            if dof not in bc_dict.keys():
                bc_dict[dof]=bc.Value
            else:
                bc_dict[dof]+=bc.Value
# key: 从1开始的自由度编号, value: 加载节点力值
eqf=np.zeros(flobal_dof_num)
for load in load_list:
    if isinstance(load,ConcentratedLoad):
        for n_label in load.get_NodesList:
            dof=get_ithdof_index(n_label,3,load.dofth)
            eqf[dof]+=load.LoadValue
    else:
        # 处理单元载荷
        for e in element_list:
            if e.label not in load.get_ElemsList:
                continue
            eqf[e.ElemDofIndexs]+=e.ProcessELoad2EqNodeForce(load)
# 找出eqF中非0值的索引和非零值,定义全局节点力字典{key:自由度索引(0-base), value:节点力}
# P_ind=eqf.nonzero()
# load_dict=dict(zip((P_ind[0]+1).tolist(),Fg[P_ind].tolist()))
load_dict={ind+1:eqf[ind] for ind in eqf.nonzero()[0] }

# -------------运行分析------------------
if bcMethod==1:
    ug,fg=Method1(KG,bc_dict,load_dict) # type: ignore
if bcMethod==2:
    ug,fg=Method2(KG,bc_dict,load_dict) # type: ignore
if bcMethod==3:
    ug,fg=Method3(KG,bc_dict,load_dict) # type: ignore
if bcMethod==4:
    ug,fg=Method4(KG,bc_dict,load_dict) # type: ignore

# -------------计算输出结果--------------------
# todo,计算外插节点应力(elem_num,8)进行,后处理中磨平
estress_array,estrain_array=np.zeros((6,element_num,element_data.shape[1]-1)),np.zeros((6,element_num,element_data.shape[1]-1))
for eind,e in enumerate(element_list):
    ue=ug[np.ix_(e.ElemDofIndexs)]
    # 插值到节点的应力,(6,8)数组
    estress=e.cal_stress_strain(ue,'stress')
    estrain=e.cal_stress_strain(ue,'strain')
    # 遍历应力分量
    estress_array[:,eind,:]=estress
    estrain_array[:,eind,:]=estrain

# 输出结果到mat文件,可以和matlab对接
result_data={
    JOB:{
        JOB_NAME:job_name,
        JOB_TIME:datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        INP_PATH:mesh_file,
    },
    MODEL:{
        # 节点信息
        NODE_DEFINE:node_data,
        NODE_VARIABLE:['U1','U2','U3'],
        NODE_NUMBER:int(node_num),
        # 单元信息
        ELEMENT_DEFINE:element_data,
        ELEMENT_TYPE:"C3D8",
        VTK_TYPE:element_list[0].VTK_TYPE,
        ELEMENT_NUMBER:int(element_num),
        ELEMENT_NODE_NUM:8,

        # 集合定义
        NSET_DEFINE:nset_data,
        ELSET_DEFINE:elset_data,
        ELEMENT_SURFACE_DEFINE:surf_data,

        # 载荷和边界
        BC_DEFINE:{ bc.name:bc.get_NodesList for bc in bc_list},
        NODE_LOAD_DEFINE:{load.name:load.get_NodesList for load in load_list if isinstance(load,ConcentratedLoad)},
        ELEMENT_LOAD_DEFINE:{load.name:load.get_ElemsList for load in load_list if not isinstance(load,ConcentratedLoad)},
        # 节点变形值
        NODE_DISPLACEMENT:np.concatenate((ug[::3].reshape(-1,1),ug[1::3].reshape(-1,1),ug[2::3].reshape(-1,1)),axis=1)
    },
    RESULT:{
        NODE_RESULT:{
            "U1":ug[::3],
            "U2":ug[1::3],
            "U3":ug[2::3],
            "Umag":np.sqrt(ug[::3]**2+ug[1::3]**2+ug[2::3]**2),
            "RF1":fg[::3],
            "RF2":fg[1::3],
            "RF3":fg[2::3],
            "RFmag":np.sqrt(fg[::3]**2+fg[1::3]**2+fg[2::3]**2),
            },
        ELEMENT_RESULT:{
            },
        INTEGRATION_POINT_RESULT:{
            "S11":estress_array[0],
            "S22":estress_array[1],
            "S33":estress_array[2],
            "S12":estress_array[3],
            "S23":estress_array[4],
            "S13":estress_array[5],
            "E11":estrain_array[0],
            "E22":estrain_array[1],
            "E33":estrain_array[2],
            "E12":estrain_array[3],
            "E23":estrain_array[4],
            "E13":estrain_array[5],
        }
    },
}

with open(export_path+job_name+'.dill','wb') as f:
    dill.dump(result_data,f)

print("结果已输出到文件:",export_path+job_name+'.dill')

